# -*- coding: utf-8 -*-

import numpy as np
import os
import SimpleITK as sitk
import pandas as pd
import multiprocessing as mp
# import torch


def file_name(file_dir):
    L = []
    path_list = os.listdir(file_dir)
    path_list.sort()  # 对读取的路径进行排序
    for filename in path_list:
        if 'nii' in filename:
            L.append(os.path.join(filename))
    return L


def computeQualityMeasures(lP, lT, label_value=1):
    quality = dict()
    labelPred = lP
    labelTrue = lT
    # labelPred = sitk.GetImagefromArray(lP, isVector=False)
    # labelTrue = sitk.GetArrayFromImage(lT, isVector=False)
    hausdorffcomputer = sitk.HausdorffDistanceImageFilter()
    hausdorffcomputer.Execute(labelTrue == label_value, labelPred == label_value)
    quality["avgHausdorff"] = hausdorffcomputer.GetAverageHausdorffDistance()
    quality["Hausdorff"] = hausdorffcomputer.GetHausdorffDistance()

    dicecomputer = sitk.LabelOverlapMeasuresImageFilter()
    dicecomputer.Execute(labelTrue == label_value, labelPred == label_value)
    quality["dice"] = dicecomputer.GetDiceCoefficient()

    return quality


seg_gt_root = 'groundtruth/'
seg_pred_root = 'predicted/'

sec_gt_root = 'mydata_section/section'
sec_pred_root = 'mydata_section/pred_section'

SEG_GT_PREFIX = 'segmentation_'
SEG_PRED_PREFIX = 'prediction_'
NIFTI_SUFFIX = '.nii.gz'

SECTION_PREFIX = 'section_'

def dice_segment(seq_list):
    dice_all = []
    avgH_all = []
    Hausdorff = []
    f = open('dice.csv','w')
    for seq_number in seq_list:
        gt_path = os.path.join(seg_gt_root, SEG_GT_PREFIX + seq_number + NIFTI_SUFFIX)
        gt = sitk.ReadImage(gt_path)
        origin = gt.GetOrigin()  # 这三句是获取的原始图像文件的位置和方向
        spacing = gt.GetSpacing()
        direction = gt.GetDirection()
        # shape = sitk.GetArrayFromImage(gt).shape
        # print(shape)
    
        pred_path = os.path.join(seg_pred_root, SEG_PRED_PREFIX + seq_number + NIFTI_SUFFIX)
        pred = sitk.ReadImage(pred_path)
        pred.SetOrigin(origin)  # 将自己的文件处理成和官方一致的位置坐标系
        pred.SetSpacing(spacing)
        pred.SetDirection(direction)
        # shape = sitk.GetArrayFromImage(pred).shape
        # print(shape)
        # sitk.WriteImage(pred, "位置路径")  # 处理完之后保存到相应的合适位置。
        quality = computeQualityMeasures(pred, gt)
        print(quality)
        f.write(seq_number+','+str(quality['dice'])+','+str(quality['avgHausdorff'])+','+str(quality['Hausdorff'])+'\n')
        dice_all.append(quality['dice'])
        avgH_all.append(quality['avgHausdorff'])
        Hausdorff.append(quality['Hausdorff'])
    f.close()
    avg_H = np.mean(avgH_all)
    hausdorff = np.mean(Hausdorff)
    dice = np.mean(dice_all)
    print('dice:',dice,'avg_H:',avg_H,'Hausdorff:',hausdorff)


def dice_section(seq_list, shared_list):
    # lists below are list(list) with the (label_value-1) as first index
    dice_all = [[] for i in range(17)]
    avgH_all = [[] for i in range(17)]
    hausdorff_all = [[] for i in range(17)]
    

    for seq_number in seq_list:
        gt_path = os.path.join(sec_gt_root, SECTION_PREFIX + seq_number + NIFTI_SUFFIX)
        pred_path = os.path.join(sec_pred_root, SECTION_PREFIX + seq_number + NIFTI_SUFFIX)
        gt = sitk.ReadImage(gt_path)
        pred = sitk.ReadImage(pred_path)
        item = {'seq_number': seq_number}
        for label_value in range(1, 18):
            quality = computeQualityMeasures(pred, gt, label_value)
            item['dice_' + str(label_value)] = quality['dice']
            item['avgHD_' + str(label_value)] = quality['avgHausdorff']
            item['HD_' + str(label_value)] = quality['Hausdorff']
            print('finished section', label_value, ', dice:', quality['dice'])
        mean_dice, mean_aHD, mean_HD = 0, 0, 0
        mean_dice_am, mean_aHD_am, mean_HD_am = 0, 0, 0
        for i in range(17):
            item_list = list(item.values())
            mean_dice += item_list[1 + i * 3]
            mean_aHD += item_list[1 + i * 3 + 1]
            mean_HD += item_list[1 + i * 3 + 2]
            if i > 5:
                mean_dice_am += item_list[1 + i * 3]
                mean_aHD_am += item_list[1 + i * 3 + 1]
                mean_HD_am += item_list[1 + i * 3 + 2]
        mean_dice /= 17
        mean_aHD /= 17
        mean_HD /= 17
        mean_dice_am /= 11
        mean_aHD_am /= 11
        mean_HD_am /= 11
        item['mean_dice'] = mean_dice
        item['mean_avgHD'] = mean_aHD
        item['mean_HD'] = mean_HD
        item['mean_dice_nobasal'] = mean_dice_am
        item['mean_avgHD_nobasal'] = mean_aHD_am
        item['mean_HD_nobasal'] = mean_HD_am
        shared_list.append(item)
        print('finished dice compute:', seq_number)
    


def get_dice_section(seq_list):
    df_list = []
    dice_section(seq_list, df_list)
    df = pd.DataFrame(df_list, columns=list(df_list[0].keys()))
    df.to_csv('dice_section.csv', index=False)

    pass


if __name__ == '__main__':
    ori_test_list = list(range(177, 197))
    ori_test_list.remove(186)

    extra_test_list = list(range(260, 345))
    extra_test_list += list(range(350, 442))

    seq_list = [str(i).zfill(5) for i in ori_test_list]
    get_dice_section(seq_list)
